"""
该脚本示例了一个 WebTransport 客户端 + CTP 交易的使用，
分为两个线程循环
- 生产者用于连接服务器并接收数据放置在队列中, 并定时处理逻辑(双协程)
- 消费者用于消耗队列中的数据, 并进行处理, 如果有必要可以将TX数据await发送给服务器
其次示例代码实现了交互的测试功能,
接收包括
- 验证信息
- 交易数据(Tick)
- 订单信息
- 仓位信息
发送包括
- 订阅Tick数据
- 多头开仓指令
- 多头平仓指令
- 空头开仓指令
- 空头平仓指令
- 撤单指令
- 查询当前仓位
- 查询历史订单
"""
import os
import dotenv
dotenv.load_dotenv()  # 从 .env 文件加载环境变量
from eztrade_client import WebTransportClient
from eztrade_client import (
    PyMessage,
    Order,
    TickData,
    OrderStatus,
    Direction,
    PositionEffect,
    Portfolio,
    OrderOperation,
    Operation,
    get_logger,
)
import dataclasses
import asyncio
import traceback
import logging
import sys
import httpx

SERVER = "127.0.0.1"
PORT = 7878
BASE_URL = f"https://{SERVER}:{PORT}"
STAGE_ID = 1
PATH = f"/wt/{STAGE_ID}"
USER_ID = 1
STAGE_PAYLOAD = {
    "id": STAGE_ID,
    "stage_type": "CTP",
    "name": f"pytest-ctp-{STAGE_ID}",
    "description": "Pytest 集成测试 CTP Stage",
    "ctp_cfg": {
        "front_td": os.getenv("CTP_FRONT_TD", "tcp://182.254.243.31:40001"),
        "front_md": os.getenv("CTP_FRONT_MD", "tcp://182.254.243.31:40011"),
        "broker_id": os.getenv("CTP_BROKER_ID", "9999"),
        "user_id": os.getenv("CTP_USER_ID", "123456"),
        "password": os.getenv("CTP_PASSWORD", "123456"),
        "auth_code": "0000000000000000",
        "app_id": os.getenv("CTP_APP_ID", "simnow_client_test"),
        "interface_type": "CTP",
    },
}

logger = get_logger("Simulate", level=logging.INFO)
logger.addFileHandler("logs/ctp_wt.log", level=logging.DEBUG)


@dataclasses.dataclass
class AccountsInfo:
    sent_orders: dict[int, Order] = dataclasses.field(default_factory=dict)
    done_orders: dict[int, Order] = dataclasses.field(default_factory=dict)
    portfolio: Portfolio | None = dataclasses.field(default=None)
    lastprice: dict[str, float] = dataclasses.field(default_factory=dict)
    once_lock: bool = False  # 仅开一次仓


ACCOUNTS_INFO = AccountsInfo()

async def process_events(client: WebTransportClient):
    """
    处理事件：持续消耗数据队列, 并根据不同类型的事件执行相应操作
    """
    while True:
        if client.buffer:
            data = client.buffer.popleft()
            if client.buffer.__len__() > 10:
                logger.warning(f"Size of buffer: {len(client.buffer)}")
            try:
                message = PyMessage.from_binary_bytes(data)
            except Exception as e:
                logger.error(f"Failed to decode message: {e}, raw: {data}")
                continue
            match message.get_data_type():
                case "Authority":  # 鉴权成功
                    auth = message.get_data()
                    logger.debug(f"UserId: {auth.client_id}")
                    logger.info("Query current positions")
                    client.send_message(PyMessage.query_account())  # 发送账户查询
                    # 需要从redis获取主力合约
                    client.send_message(PyMessage.subscribe_tick("au2511"))
                    # client.send_message(PyMessage.subscribe_tick("cu2511"))
                case "TickData":
                    tickdata: TickData = message.get_data()
                    ACCOUNTS_INFO.lastprice[tickdata.symbol] = tickdata.last
                    logger.info(f"Received TickData: {tickdata}")
                    # ---------------------------------------------
                    if not ACCOUNTS_INFO.once_lock:
                        import random
                        direction = random.choice([Direction.Long, Direction.Short])
                        position_effect = random.choice([PositionEffect.Open, PositionEffect.Close])
                        order = Order(tickdata.symbol, direction, position_effect, 1) # 开1手(需要考虑乘数)
                        msg = PyMessage(order)
                        client.send_message(msg)
                        logger.debug(f"Send Order: {order}")
                case "Order":
                    order: Order = message.get_data()
                    # 处理订单数据
                    if order.id not in ACCOUNTS_INFO.sent_orders:
                        logger.debug(f"Order {order} has been sent.")
                        ACCOUNTS_INFO.once_lock = True
                        ACCOUNTS_INFO.sent_orders.update({order.id: order})
                        if (
                            ACCOUNTS_INFO.portfolio
                            and ACCOUNTS_INFO.portfolio.balance_available / 100
                            < 5_000_000
                        ):
                            # 撤单指令
                            msg = PyMessage(OrderOperation(order.id, Operation.Cancel))
                            logger.debug(f"Cancel Order: {order.id}")
                            client.send_message(msg)
                    if order.id in ACCOUNTS_INFO.sent_orders:  # 完成时事件回传
                        if order.status in [
                            OrderStatus.Filled,
                            OrderStatus.PartiallyFilled,
                            OrderStatus.Canceled,
                            OrderStatus.Rejected,
                        ]:
                            logger.info(f"Queried Order has done! {order}")
                            ACCOUNTS_INFO.sent_orders.pop(order.id)
                            ACCOUNTS_INFO.done_orders.update({order.id: order})
                case "Portfolio":
                    portfolio: Portfolio = message.get_data()
                    logger.info(f"Got portfolio information: {portfolio}")
                    logger.info(f"Current positions: {portfolio.positions}")
                    ACCOUNTS_INFO.portfolio = portfolio
                case "String":
                    data: str = message.get_data()
                    logger.info(f"Received string data: {data}")
                case _:
                    logger.info(f"Received data: {message}")
        else:
            await asyncio.sleep(0.001)  # 如果没有数据，交出控制权

        if client.stop_flag.is_set():
            logger.info("客户端接收停止标志已设置, 退出消费者循环")
            break


async def simulate_query(client: WebTransportClient):
    """
    模拟查询：发送查询指令, 定时更新仓位信息
    """
    while True:
        if client.stop_flag.is_set():
            logger.info("客户端控制停止标志已设置, 退出控制循环")
            break
        if not ACCOUNTS_INFO.portfolio:
            # 还未初始化
            await asyncio.sleep(1)
            continue
        # 查询持仓信息
        msg = PyMessage.query_account()
        client.send_message(msg)  # 发送账户查询
        if ACCOUNTS_INFO.portfolio.positions:
            # 打印持仓信息
            value_all = 0
            logger.info(
                f"current balance: {ACCOUNTS_INFO.portfolio.balance_available / 100}"
            )  # 仅可用部分
            logger.info(
                f"current b_value: {ACCOUNTS_INFO.portfolio.balance / 100}"
            )
            value_all += (
                ACCOUNTS_INFO.portfolio.balance / 100
            )  # (包括做空获得的不可用部分)
            for position in ACCOUNTS_INFO.portfolio.positions:
                last = ACCOUNTS_INFO.lastprice.get(position.symbol, 0)
                logger.info(f"Position: {position}, value={last * position.quantity}")
                if position.side == Direction.Long:
                    value_all += last * position.quantity
                elif position.side == Direction.Short:
                    value_all -= last * position.quantity
            logger.info(f"Total portfolio value: {value_all}")  # 总值
        await asyncio.sleep(5)  # 控制逻辑间隔


async def my_strategy(client: WebTransportClient):
    pass

async def main():
    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.post("/stages/create", json=STAGE_PAYLOAD)
        assert resp.status_code in [409, 201], f"Create stage failed: {resp.text}"
    await asyncio.sleep(1) # 等待CTP登录

    client = WebTransportClient()
    client.create(SERVER, PORT, PATH, USER_ID)
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(process_events(client))
            tg.create_task(simulate_query(client))
    except ExceptionGroup as es:
        for e in es.exceptions:
            logger.error(
                f"主循环异常退出: {e}: {''.join(traceback.format_exception(e))} "
            )
    except asyncio.CancelledError:
        logger.info("主循环取消退出")
    finally:
        client.stop_loop_thread()

    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.delete(f"stages/{STAGE_ID}")
        if resp.status_code != 204:
            logger.error(f"Delete stage failed: {resp.text}")
        else:
            logger.info(f"Stage {STAGE_ID} deleted successfully.")
    sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())
